"""
翻译质量评估COMET: https://github.com/Unbabel/COMET
提供 0 到 1 的分数，其中 1 表示完美的翻译

COMET是为了学习预测人类对MT质量的判断而设计的。它通过使用神经系统首先将MT生成的译文，
参考译文和源语言文本映射到神经意义表征中来实现这一目的。然后，它利用这些表示来学习预测质量分数，
该质量分数被显式优化以与人类对翻译质量的判断相关联。

可解释 COMET（XCOMET）： Unbabel/XCOMET-XXL - 我们最新的模型被训练以识别错误范围并分配最终质量分数，
从而产生一个可解释的神经指标。我们提供 XXL 版本，包含 107 亿参数，以及 XL 版本，
包含 35 亿参数（ Unbabel/XCOMET-XL ）。这些模型与 MQM 的最高相关性，是我们表现最好的评估模型。

支持东盟语的翻译质量评估, https://hf-mirror.com/Unbabel/wmt22-comet-da

无参考评估：comet-score -s src.txt -t hyp1.txt --model Unbabel/wmt22-cometkiwi-da

"""
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from comet import download_model, load_from_checkpoint
from threading import RLock
from typing import List, Dict


class CometEvaluator(object):

    lock = RLock()

    def __new__(cls, *args, **kwargs):
        with CometEvaluator.lock:
            if not hasattr(CometEvaluator, "_instance"):
                CometEvaluator._instance = object.__new__(cls)
                repo_id = "Unbabel/wmt22-comet-da"
                # repo_id = "Unbabel/XCOMET-XL"
                model_path = download_model(repo_id)
                cls.model = load_from_checkpoint(model_path)
        return CometEvaluator._instance

    def evaluate(self, data: List[Dict]):
        report = self.model.predict(data, batch_size=8, gpus=1)
        return report


if __name__ == '__main__':
    data = [
        {
            "src": "Dem Feuer konnte Einhalt geboten werden",
            "mt": "The fire could be stopped",
            "ref": "They were able to control the fire."
        },
        {
            "src": "Schulen und Kindergärten wurden eröffnet.",
            "mt": "Schools and kindergartens were open",
            "ref": "Schools and kindergartens opened"
        }
    ]
    comet = CometEvaluator()
    report = comet.evaluate(data)
    print(report)
    print(report.scores)
    print(report.system_score)
